from zhipuai import ZhipuAI
from config.config_manager import ConfigManager


class ZhipuAIHandler:
    def __init__(self):
        """
        初始化智谱AI处理器，从配置文件加载API Key
        """
        self.config_manager = ConfigManager()
        self.config = self.config_manager.load_config()
        api_key = self.config.get("AI_ZP_ApiKey", "").strip()

        if not api_key:
            raise ValueError("配置文件中未找到有效的 ZhiPuApiKey")

        try:
            self.client = ZhipuAI(api_key=api_key)
        except Exception as e:
            raise Exception(f"初始化 ZhipuAI 客户端时出错: {e}")

    def generate_strategy(self, table_script: str) -> dict:

        if not hasattr(self, 'client') or self.client is None:
            raise Exception("ZhipuAI 客户端未正确初始化")

        try:
            response = self.client.chat.completions.create(
                model="glm-4-plus",
                messages=[
                    {"role": "user", "content": "作为一名测试专家，你通晓各种信息项的最佳生成策略"},
                    {"role": "assistant",
                     "content": "当然，我擅长使用 python中的 faker库来为各种信息项制定模拟数据的生成策略"},
                    {"role": "user",
                     "content": "好的，基于下面的数据库表的建表语句，为其中的各个字段制定模拟数据的生成策略："},
                    {"role": "user", "content": table_script},
                    {"role": "user",
                     "content": "在生成之前，先检测字段的数据类型，如果不是文本类型，则使用随机的生成策略，比如近 2 年的随机时间，随机数值等。"},
                    {"role": "user",
                     "content": "如果是文本类型，则基于字段的注释，分析其语义，基于语义与 faker库制定最佳的文本生成策略，切记要使用faker库中存在的生成策略，不可胡编乱造！"},
                    {"role": "user", "content": "以json格式返回，key值是字段名称，value值是一个对象，包括基于faker库的可执行的生成策略(不要包含lambda,faker库的名称固定为faker)，标签名称为 strategy,包括字段的数据类型，标签名称为 data_type，包括字段的注释，标签名称为 comment，包括字段是否生成测试数据的标志位，标签名称为 is_gen"},
                    {"role": "user",
                     "content": "最后在返回的信息中，我只要上面提及的那个json格式，由{}包裹的 json数据，其它任何信息或字符都不要！"}
                ]
            )
            return response.choices[0].message.content
        except Exception as e:
            raise Exception(f"生成策略时出错: {e}")
